import csv
import ast
import matplotlib.pyplot as plt
import numpy as np

def read_extras_csv(file_path="extras_log_300Hz.csv"):
    keys = [
        "ankle_angle",
        "ankle_torques",
        "left_imu_lin_acc",
        "left_imu_ang_vel",
        "right_imu_lin_acc",
        "right_imu_ang_vel",
        "left_ankle_hight",
        "right_ankle_hight",
        "joint_vel"
    ]

    data = {k: [] for k in keys}

    with open(file_path, "r", newline="") as f:
        reader = csv.reader(f)
        for row_idx, row in enumerate(reader):
            print(f"Record {row_idx+1}:")
            for i, k in enumerate(keys):
                vec = ast.literal_eval(row[i])
                data[k].append(vec)
                print(f"  {k} length: {len(vec)}")
    
    # 转成 NumPy 数组，每个 key 是二维数组 (行数, 向量长度)
    for k in keys:
        data[k] = np.vstack(data[k])

    return data

def scatter_extras_data(data):
    for key, arr in data.items():
        plt.figure(figsize=(10,4))
        time_steps = np.arange(arr.shape[0])[:, None]  # 每条记录的时间步
        for i in range(arr.shape[1]):
            plt.scatter(time_steps, arr[:, i], s=5, label=f"{key}_{i}" if arr.shape[1] <= 10 else None)
        plt.title(f"Scatter plot of {key}")
        plt.xlabel("Time step")
        plt.ylabel(key)
        if arr.shape[1] <= 10:
            plt.legend()
        plt.tight_layout()
        plt.show()

if __name__ == "__main__":
    extras_data = read_extras_csv("extras_log_300Hz.csv")
    scatter_extras_data(extras_data)
